import csv
import os
import numpy as np
from scipy.fftpack import fft
import matplotlib.pyplot as plt

###本文件用于
def FFT(Fs, data):
    """
    对输入信号进行FFT
    :param Fs:  采样频率
    :param data:待FFT的序列
    :return:
    """
    L = len(data)  # 信号长度
    N = np.power(2, np.ceil(np.log2(L)))  # 下一个最近二次幂，也即N个点的FFT
    result = np.abs(fft(x=data, n=int(N))) / L * 2  # N点FFT
    axisFreq = np.arange(int(N / 2)) * Fs / N  # 频率坐标
    result = result[range(int(N / 2))]  # 因为图形对称，所以取一半
    return axisFreq, result

def main():
    Fs = 2560
    path = ['../../baseline/LSTM/learn_files','../baseline/LSTM/test_files']
    for p in path:
        p = p + '/feature'
        files = os.listdir(p)
        for name in sorted(files):
            f_path = p + '/' + name
            save_data = []
            with open(f_path) as f:
                reader = csv.reader(f)
                for line in reader:
                    line = [float(x) for x in line]
                    print(len(line))
                    x, result = FFT(Fs, line)
                    result = result.tolist()
                    print(len(x))
                    print(len(result))
                    input()
                    save_data.append(result)
                    t = np.linspace(0, 1, Fs)
                    # 绘图
                    fig1 = plt.figure(figsize=(16, 9))
                    plt.title('original data')
                    plt.plot(t, line)
                    plt.xlabel('time/s')
                    plt.ylabel('Amplitude')
                    plt.grid()
                    plt.show()
                    fig2 = plt.figure(figsize=(16, 9))
                    plt.title('FFT')
                    plt.plot(x, result)
                    plt.xlabel('Frequency/Hz')
                    plt.ylabel('Amplitude')
                    plt.grid()
                    plt.show()
                    input()
            with open(name,'w',newline='') as f:
                writer = csv.writer((f))
                writer.writerows(save_data)


if __name__ == '__main__':
    main()